package com.hzya.frame;

import net.sf.jsqlparser.expression.Expression;
import net.sf.jsqlparser.parser.CCJSqlParserUtil;
import net.sf.jsqlparser.statement.select.PlainSelect;
import net.sf.jsqlparser.statement.select.Select;
import org.apache.nifi.annotation.behavior.InputRequirement;
import org.apache.nifi.annotation.behavior.InputRequirement.Requirement;
import org.apache.nifi.annotation.behavior.Stateful;
import org.apache.nifi.annotation.behavior.WritesAttribute;
import org.apache.nifi.annotation.behavior.WritesAttributes;
import org.apache.nifi.annotation.documentation.CapabilityDescription;
import org.apache.nifi.annotation.documentation.Tags;
import org.apache.nifi.components.PropertyDescriptor;
import org.apache.nifi.components.state.Scope;
import org.apache.nifi.components.state.StateManager;
import org.apache.nifi.components.state.StateMap;
import org.apache.nifi.dbcp.DBCPService;
import org.apache.nifi.flowfile.FlowFile;
import org.apache.nifi.processor.AbstractProcessor;
import org.apache.nifi.processor.ProcessContext;
import org.apache.nifi.processor.ProcessSession;
import org.apache.nifi.processor.Relationship;
import org.apache.nifi.processor.exception.ProcessException;
import org.apache.nifi.processor.util.StandardValidators;

import java.nio.charset.StandardCharsets;
import java.sql.Connection;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.sql.Statement;
import java.util.*;

@Tags({"sql", "database", "pagination", "custom", "jdbc", "oracle", "mysql", "sqlserver"})
@CapabilityDescription("基于基本查询生成分页SQL查询，并有状态地跟踪列的最大值，以仅获取新数据。每个输出FlowFile的内容都是特定页面的分页SQL查询。支持Oracle、MySQL和SQL Server。")
@InputRequirement(Requirement.INPUT_REQUIRED)
@Stateful(scopes = {Scope.CLUSTER}, description = "为每个表名存储指定列的最大值。这允许处理器在后续运行中仅查询新行。")
@WritesAttributes({@WritesAttribute(attribute = "page", description = "生成的SQL查询的当前页码。"), @WritesAttribute(attribute = "startRow", description = "此页的起始行号。"), @WritesAttribute(attribute = "endRow", description = "此页的结束行号。")})
public class DevGeneratePaginatedSqlProcessor extends AbstractProcessor {
    // 属性定义
    public static final PropertyDescriptor DBCP_SERVICE = new PropertyDescriptor.Builder().name("Database Connection Pooling Service").description("提供与数据库连接的控制器服务。").required(true).identifiesControllerService(DBCPService.class).build();

    public static final PropertyDescriptor BASE_SQL_QUERY = new PropertyDescriptor.Builder().name("Base SQL Query").description("要分页的基本SQL查询。应该是SELECT语句。该值将从传入FlowFile的“executesql”属性中读取。").required(true).expressionLanguageSupported(true).addValidator(StandardValidators.NON_EMPTY_VALIDATOR).build();

    public static final PropertyDescriptor TABLE_NAME = new PropertyDescriptor.Builder().name("Table Name for State").description("正在查询的表的唯一名称。这被用作存储状态（最大值）的键。该值将从“queryTableName”属性中读取。").required(true).expressionLanguageSupported(true).addValidator(StandardValidators.NON_EMPTY_VALIDATOR).build();

    public static final PropertyDescriptor MAX_VALUE_COLUMN = new PropertyDescriptor.Builder().name("Maximum-Value Column").description("将跟踪其最大值以获取增量数据的列的名称（例如，时间戳或自动递增的ID）。将从“maximumValueColumns”属性读取该值。").required(true).expressionLanguageSupported(true).addValidator(StandardValidators.NON_EMPTY_VALIDATOR).build();

    public static final PropertyDescriptor PAGE_SIZE = new PropertyDescriptor.Builder().name("Page Size").description("每页中要包含的行数。该值将从“pagerow”属性读取，默认为1000。").defaultValue("1000").required(true).expressionLanguageSupported(true).addValidator(StandardValidators.POSITIVE_INTEGER_VALIDATOR).build();

    // 关系定义
    public static final Relationship REL_SUCCESS = new Relationship.Builder().name("success").description("分页SQL查询已成功生成并发送到此关系。").build();

    public static final Relationship REL_FAILURE = new Relationship.Builder().name("failure").description("如果在处理过程中发生错误，则将原始FlowFile路由到此关系。").build();

    private final List<PropertyDescriptor> descriptors;
    private final Set<Relationship> relationships;

    public DevGeneratePaginatedSqlProcessor() {
        final List<PropertyDescriptor> descriptors = new ArrayList<>();
        descriptors.add(DBCP_SERVICE);
        descriptors.add(BASE_SQL_QUERY);
        descriptors.add(TABLE_NAME);
        descriptors.add(MAX_VALUE_COLUMN);
        descriptors.add(PAGE_SIZE);
        this.descriptors = Collections.unmodifiableList(descriptors);

        final Set<Relationship> relationships = new HashSet<>();
        relationships.add(REL_SUCCESS);
        relationships.add(REL_FAILURE);
        this.relationships = Collections.unmodifiableSet(relationships);
    }

    @Override
    public Set<Relationship> getRelationships() {
        return this.relationships;
    }

    @Override
    public final List<PropertyDescriptor> getSupportedPropertyDescriptors() {
        return this.descriptors;
    }

    @Override
    public void onTrigger(ProcessContext context, ProcessSession session) throws ProcessException {
        FlowFile flowFile = session.get();
        if (flowFile == null) {
            return;
        }

        // 获取属性值
        final String baseSql = context.getProperty(BASE_SQL_QUERY).evaluateAttributeExpressions(flowFile).getValue();
        final String tableName = context.getProperty(TABLE_NAME).evaluateAttributeExpressions(flowFile).getValue();
        final String maxValColumn = context.getProperty(MAX_VALUE_COLUMN).evaluateAttributeExpressions(flowFile).getValue();
        final Integer pageSize = context.getProperty(PAGE_SIZE).evaluateAttributeExpressions(flowFile).asInteger();

        // 验证属性
        if (baseSql == null || tableName == null || maxValColumn == null || pageSize == null) {
            getLogger().error("缺少必需的属性。路由失败");
            session.transfer(flowFile, REL_FAILURE);
            return;
        }

        final DBCPService dbcpService = context.getProperty(DBCP_SERVICE).asControllerService(DBCPService.class);
        final StateManager stateManager = context.getStateManager();
        final Scope scope = Scope.CLUSTER;

        try {
            final StateMap stateMap = stateManager.getState(scope);
            final String oldMaximumValue = stateMap.get(tableName);

            String queriedMaximumValue = null;
            long rowNumber = 0L;

            try (Connection conn = dbcpService.getConnection(); Statement stmt = conn.createStatement()) {
                String dbType = conn.getMetaData().getDatabaseProductName().toLowerCase();

                //增量同步
                String whereClause = "";
                if (oldMaximumValue != null && !oldMaximumValue.isEmpty()) {
                    whereClause = " WHERE " + maxValColumn + " > '" + oldMaximumValue + "'";
                }

                // 执行 COUNT 和 MAX 查询
                String countMaxSql = "";
                if ("oracle".equals(dbType) || "mysql".equals(dbType)) {
                    countMaxSql = "SELECT COUNT(*) AS rowNumber, MAX(" + maxValColumn + ") AS maximumValue FROM (" + baseSql + ") t " + whereClause;
                } else if ("microsoft sql server".equals(dbType)) {
                    StringBuffer montageSql = new StringBuffer();
                    montageSql.append("WITH SortedData AS (");

                    if (!"".equals(whereClause)) {
//                        if (baseSql.toUpperCase().contains("WHERE")) {
//                            //替换为and
//                            whereClause = whereClause.replace("WHERE", "AND");
//                        }
                        //where条件不能直接加载原始sql语句后面，否则order by xx asc where xxx，导致语法问题，而是要加载表名后面
//                        montageSql.append(whereClause);
                        whereClause = whereClause.replace("WHERE", "");
//                        whereClause = whereClause.replace("AND", "");
                        montageSql.append(insertWhereClause(baseSql, whereClause));
                    } else {
                        montageSql.append(baseSql);
                    }
                    if (baseSql.toUpperCase().contains("ORDER BY")) {
                        //确保排序生效
                        montageSql.append(" OFFSET 0 ROWS");
                    }
                    montageSql.append(")");
                    montageSql.append("\n");
                    montageSql.append("SELECT");
                    montageSql.append("    COUNT(*) AS rowNumber,");
                    montageSql.append("    MAX(ts) AS maximumValue");
                    montageSql.append("    FROM SortedData");

                    countMaxSql = montageSql.toString();
                } else {
                    throw new ProcessException("不支持的数据库类型: " + dbType);
                }

                getLogger().info("正在执行计数/最大值查询： {}", countMaxSql);
                ResultSet rs = stmt.executeQuery(countMaxSql);

                if (rs.next()) {
                    rowNumber = rs.getLong("rowNumber");
                    queriedMaximumValue = rs.getString("maximumValue");
                }
                rs.close();

                getLogger().info("表“{}”的查询结果：新行总数={}，新的最大值={}", tableName, rowNumber, queriedMaximumValue);

                // 无新数据时删除 FlowFile 并返回
                if (rowNumber == 0 || queriedMaximumValue == null) {
                    getLogger().info("找不到表“{}”的新数据。正在删除原始FlowFile。", tableName);
                    session.remove(flowFile);
                    return;
                }

                // 生成分页 SQL
                long totalPages = (rowNumber + pageSize - 1) / pageSize;

                for (int page = 0; page < totalPages; page++) {
                    long startRow = (long) page * pageSize + 1;
                    long endRow = Math.min((long) (page + 1) * pageSize, rowNumber);

                    //全量同步
                    if ("".equals(whereClause)) {
                        whereClause = " WHERE " + maxValColumn + " <= '" + queriedMaximumValue + "'";
                    }

                    String paginatedSql;
                    if ("oracle".equals(dbType)) {
                        paginatedSql = "SELECT * FROM (\n" + "    SELECT inner_query.*, ROWNUM rn FROM (\n" + "        " + "select * from (" + baseSql + ")" + whereClause + "\n" + "    ) inner_query WHERE ROWNUM <= " + endRow + "\n" + ") WHERE rn >= " + startRow;
                    } else if ("mysql".equals(dbType)) {
                        long offset = startRow - 1;
                        paginatedSql = "select * from (" + baseSql + ")" + whereClause + " LIMIT " + pageSize + " OFFSET " + offset;
                    } else if ("microsoft sql server".equals(dbType)) {

                        whereClause = whereClause.replace("WHERE", "");
//                        whereClause = whereClause.replace("AND", "");

                        StringBuffer splicingSql = new StringBuffer();
                        splicingSql.append(baseSql);

                        splicingSql.append(" OFFSET ");
                        splicingSql.append((startRow - 1));
                        splicingSql.append(" ROWS FETCH NEXT ");
                        splicingSql.append(pageSize);
                        splicingSql.append(" ROWS ONLY");

                        paginatedSql = insertWhereClause(splicingSql.toString(), whereClause);
                    } else {
                        throw new ProcessException("不支持的数据库类型: " + dbType);
                    }

                    FlowFile newFlowFile = session.create(flowFile);
                    newFlowFile = session.write(newFlowFile, out -> out.write(paginatedSql.getBytes(StandardCharsets.UTF_8)));

                    Map<String, String> attributes = new HashMap<>();
                    attributes.put("page", String.valueOf(page + 1));
                    attributes.put("startRow", String.valueOf(startRow));
                    attributes.put("endRow", String.valueOf(endRow));
                    newFlowFile = session.putAllAttributes(newFlowFile, attributes);

                    session.transfer(newFlowFile, REL_SUCCESS);
                }

                // 更新状态
                final Map<String, String> newState = new HashMap<>(stateMap.toMap());
                newState.put(tableName, queriedMaximumValue);
                stateManager.setState(newState, scope);

                // 删除原始 FlowFile
                session.remove(flowFile);
            } catch (SQLException e) {
                getLogger().error("由于{}，无法执行SQL查询；故障路由", e.getMessage(), e);
                session.transfer(flowFile, REL_FAILURE);
            }
        } catch (Exception e) {
            getLogger().error("处理过程中出错；故障路由", e);
            session.transfer(flowFile, REL_FAILURE);
        }
    }

    public static String insertWhereClause(String originalSql, String additionalCondition) throws Exception {
        // 解析SQL
        net.sf.jsqlparser.statement.Statement statement = CCJSqlParserUtil.parse(originalSql);
        if (!(statement instanceof Select)) {
            throw new IllegalArgumentException("仅支持SELECT语句");
        }

        Select select = (Select) statement;
        PlainSelect plainSelect = (PlainSelect) select.getSelectBody();

        // 构建新WHERE条件
        Expression newWhere = buildNewWhere(plainSelect.getWhere(), additionalCondition);

        // 更新WHERE条件
        plainSelect.setWhere(newWhere);

        // 重新生成SQL字符串
        return select.toString();
    }

    private static Expression buildNewWhere(Expression existingWhere, String additionalCondition) throws Exception {
        if (existingWhere == null) {
            // 无原始WHERE，直接使用新条件
            return CCJSqlParserUtil.parseCondExpression(additionalCondition);
        }

        // 合并新旧条件 (AND连接)
        Expression newCondition = CCJSqlParserUtil.parseCondExpression(additionalCondition);
        return CCJSqlParserUtil.parseCondExpression(existingWhere + " AND " + newCondition);
    }

//    public static void main(String[] args) throws Exception {
//        String originalSql = "SELECT * FROM so_sale ORDER BY ts, csaleid ASC OFFSET 0 ROWS FETCH NEXT 10 ROWS ONLY";
//        String additionalCondition = "status = 'ACTIVE' AND amount > 100";
//
//        String modifiedSql = insertWhereClause(originalSql, additionalCondition);
//        System.out.println("Modified SQL:\n" + modifiedSql);
//    }

}